#!/usr/bin/env python
# -*- coding: UTF-8 -*-
# @Project:     python_demo
# @File:        u_db.py
# @Author:      Fan GuiChuan
# @Date:        2024/7/3 16:09
# @Version:     Python3.7
# @Description: 数据库工具类
import threading

from playhouse.pool import PooledPostgresqlDatabase, PooledSqliteDatabase, PooledMySQLDatabase
from peewee import OperationalError, InterfaceError
from playhouse.shortcuts import ReconnectMixin


class ReconnectPooledPostgresqlDatabase(ReconnectMixin, PooledPostgresqlDatabase):
    """
    支持重连机制的数据库连接池类，可以通过扩展reconnect_errors来支持重连场景
    """
    reconnect_errors = (
        # Postgres error examples:
        (OperationalError, 'terminat'),
        (InterfaceError, 'connection already closed')
    )


class ReconnectPooledMySQLDatabase(ReconnectMixin, PooledMySQLDatabase):
    pass


class ReconnectPooledSqliteDatabase(ReconnectMixin, PooledSqliteDatabase):
    pass


class PostgresManager(object):
    """
    PG数据库管理类
    """
    __instance_lock = threading.Lock()
    __database = None

    @classmethod
    def get_database(cls, **kwargs):
        if cls.__database is None:
            with cls.__instance_lock:
                cls.__database = PooledPostgresqlDatabase(
                    database=kwargs.get("database"),
                    host=kwargs.get("host", "127.0.0.1"),
                    port=int(kwargs.get("port", 5432)),
                    user=kwargs.get("user", "postgres"),
                    password=kwargs.get("password", "postgres"),
                    max_connections=int(kwargs.get("max_connections", 50)),
                    stale_timeout=int(kwargs.get("stale_timeout", 600))
                )
        return cls.__database


class MysqlManager(object):
    """
    Mysql数据库管理类
    """
    __instance_lock = threading.Lock()
    __database = None

    @classmethod
    def get_database(cls, **kwargs):
        if cls.__database is None:
            with cls.__instance_lock:
                cls.__database = ReconnectPooledMySQLDatabase(
                    database=kwargs.get("database"),
                    host=kwargs.get("host", "127.0.0.1"),
                    port=int(kwargs.get("port", 3366)),
                    user=kwargs.get("user", "root"),
                    password=kwargs.get("password", "123456"),
                    max_connections=int(kwargs.get("max_connections", 50)),
                    stale_timeout=int(kwargs.get("stale_timeout", 600))
                )
        return cls.__database


class SqliteManager(object):
    """
    PG数据库管理类
    """
    __instance_lock = threading.Lock()
    __database = None

    @classmethod
    def get_database(cls, **kwargs):
        if cls.__database is None:
            with cls.__instance_lock:
                cls.__database = ReconnectPooledSqliteDatabase(
                    database=kwargs.get("database"),
                    max_connections=int(kwargs.get("max_connections", 50)),
                    pragmas={'journal_mode': 'wal', 'cache_size': -1024 * 64}
                )
        return cls.__database
